#####################################
######      Milvus知识库管理    #######
#####################################
import os
import gradio as gr
import asyncio
import re

def ensure_event_loop():
    try:
        asyncio.get_running_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)

def normalize_collection_name(name):
    """规范化集合名称，确保符合Milvus命名规则"""
    # 将连字符、空格等特殊字符替换为下划线
    normalized = re.sub(r'[^a-zA-Z0-9_]', '_', name)
    # 确保以字母或下划线开头
    if normalized and not normalized[0].isalpha() and normalized[0] != '_':
        normalized = 'kb_' + normalized
    # 如果为空，使用默认名称
    if not normalized:
        normalized = 'knowledge_base'
    # 限制长度（Milvus集合名称长度限制）
    if len(normalized) > 64:
        normalized = normalized[:64]
    return normalized

from llama_index.core import VectorStoreIndex, Settings, SimpleDirectoryReader
from llama_index.core.schema import TextNode
from llama_index.vector_stores.milvus import MilvusVectorStore
from llama_index.embeddings.dashscope import (
    DashScopeEmbedding,
    DashScopeTextEmbeddingModels,
    DashScopeTextEmbeddingType,
)
from upload_file import *

# Milvus配置
MILVUS_HOST = os.getenv("MILVUS_HOST", "localhost")
MILVUS_PORT = os.getenv("MILVUS_PORT", "19530")
MILVUS_USER = os.getenv("MILVUS_USER", "")
MILVUS_PASSWORD = os.getenv("MILVUS_PASSWORD", "")

STRUCTURED_FILE_PATH = "File/Structured"
UNSTRUCTURED_FILE_PATH = "File/Unstructured"
TMP_NAME = "tmp_abcd"

EMBED_MODEL = DashScopeEmbedding(
    model_name=DashScopeTextEmbeddingModels.TEXT_EMBEDDING_V2,
    text_type=DashScopeTextEmbeddingType.TEXT_TYPE_DOCUMENT,
)

# 设置嵌入模型
Settings.embed_model = EMBED_MODEL

def get_milvus_collections():
    ensure_event_loop()
    """获取Milvus中所有集合名称"""
    try:
        from pymilvus import connections, utility
        
        # 连接到Milvus
        connections.connect(
            alias="default",
            host=MILVUS_HOST,
            port=MILVUS_PORT,
            user=MILVUS_USER if MILVUS_USER else None,
            password=MILVUS_PASSWORD if MILVUS_PASSWORD else None
        )
        
        # 获取所有集合
        collections = utility.list_collections()
        print(f"Milvus中的所有集合: {collections}")
        connections.disconnect("default")
        return collections
    except Exception as e:
        print(f"获取Milvus集合失败: {e}")
        return []

def create_milvus_vector_store(collection_name, dimension=1536):
    ensure_event_loop()
    """创建Milvus向量存储"""
    try:
        # 规范化集合名称
        normalized_name = normalize_collection_name(collection_name)
        print(f"原始集合名称: {collection_name}, 规范化后: {normalized_name}")
        
        # 先测试嵌入模型获取正确的维度
        try:
            test_embedding = EMBED_MODEL.get_text_embedding("测试文本")
            actual_dimension = len(test_embedding)
            print(f"嵌入模型实际维度: {actual_dimension}")
            dimension = actual_dimension
        except Exception as embed_error:
            print(f"获取嵌入维度失败，使用默认维度 {dimension}: {embed_error}")
        
        vector_store = MilvusVectorStore(
            collection_name=normalized_name,
            host=MILVUS_HOST,
            port=MILVUS_PORT,
            user=MILVUS_USER if MILVUS_USER else None,
            password=MILVUS_PASSWORD if MILVUS_PASSWORD else None,
            dim=dimension,
            overwrite=True
        )
        return vector_store
    except Exception as e:
        print(f"创建Milvus向量存储失败: {e}")
        import traceback
        traceback.print_exc()
        return None

def create_milvus_collection(collection_name, dimension=1536):
    ensure_event_loop()
    """创建Milvus集合（原生pymilvus方式）"""
    try:
        # 规范化集合名称
        normalized_name = normalize_collection_name(collection_name)
        print(f"原始集合名称: {collection_name}, 规范化后: {normalized_name}")
        
        from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType, utility
        
        # 连接到Milvus
        connections.connect(
            alias="default",
            host=MILVUS_HOST,
            port=MILVUS_PORT,
            user=MILVUS_USER if MILVUS_USER else None,
            password=MILVUS_PASSWORD if MILVUS_PASSWORD else None
        )
        
        # 清理可能存在的旧集合
        if utility.has_collection(collection_name):
            utility.drop_collection(collection_name)
            print(f"已删除旧集合: {collection_name}")
        if utility.has_collection(normalized_name):
            utility.drop_collection(normalized_name)
            print(f"已删除旧集合: {normalized_name}")
        
        # 创建集合schema
        fields = [
            FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
            FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535),
            FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=dimension)
        ]
        schema = CollectionSchema(fields, description=f"知识库: {collection_name}")
        
        # 创建集合
        collection = Collection(collection_name, schema)
        print(f"✅ 集合 {collection_name} 创建成功")
        
        # 创建索引
        index_params = {
            "metric_type": "COSINE",
            "index_type": "IVF_FLAT",
            "params": {"nlist": 128}
        }
        collection.create_index(field_name="embedding", index_params=index_params)
        print("✅ 索引创建成功")
        
        connections.disconnect("default")
        return collection
        
    except Exception as e:
        print(f"❌ 创建Milvus集合失败: {e}")
        import traceback
        traceback.print_exc()
        return None

def create_unstructured_milvus_db(db_name: str, label_name: list):
    ensure_event_loop()
    print(f"Milvus知识库名称为：{db_name}，类目名称为：{label_name}")

    if label_name is None:
        gr.Info("没有选择类目")
        return
    elif len(db_name) == 0:
        gr.Info("没有命名知识库")
        return

    # 检查Milvus集合是否已存在
    existing_collections = get_milvus_collections()
    print(f"现有Milvus集合: {existing_collections}")

    normalized_name = normalize_collection_name(db_name)
    if db_name in existing_collections or normalized_name in existing_collections:
        gr.Info("Milvus知识库已存在，请换个名字或删除原来知识库再创建")
        return

    try:
        gr.Info("正在手动创建Milvus知识库...")

        # 加载文档
        documents = []
        for label in label_name:
            label_path = os.path.join(UNSTRUCTURED_FILE_PATH, label)
            print(f"检查路径: {label_path}, 存在: {os.path.exists(label_path)}")
            if os.path.exists(label_path):
                docs = SimpleDirectoryReader(label_path).load_data()
                print(f"从 {label} 加载了 {len(docs)} 个文档")
                documents.extend(docs)

        print(f"总共加载了 {len(documents)} 个文档")
        if not documents:
            gr.Info("没有找到文档数据")
            return

        # 创建Milvus向量存储
        print(f"开始创建Milvus向量存储，集合名称: {db_name}")
        
        # 使用抽取的方法创建Milvus集合
        collection = create_milvus_collection(db_name)
        if collection is None:
            gr.Info("创建Milvus集合失败")
            return

        print("Milvus向量存储创建成功，开始手动处理文档...")

        from llama_index.core.schema import TextNode
        from llama_index.core.node_parser import SentenceSplitter

        text_splitter = SentenceSplitter(chunk_size=512, chunk_overlap=50)
        all_texts = []
        all_embeddings = []
        
        for doc in documents:
            text_chunks = text_splitter.split_text(doc.get_content())
            print(f"文档 {doc.get_doc_id()} 分割为 {len(text_chunks)} 个块")
            for i, chunk in enumerate(text_chunks):
                if chunk.strip() and len(chunk.strip()) >= 10:
                    try:
                        embedding = EMBED_MODEL.get_text_embedding(chunk)
                        if embedding is None or len(embedding) == 0:
                            print(f"警告：节点 {i} 的嵌入向量为空，跳过")
                            continue
                        print(f"节点 {i} 嵌入向量维度: {len(embedding)}")
                        all_texts.append(chunk)
                        all_embeddings.append(embedding)
                    except Exception as embed_error:
                        print(f"警告：节点 {i} 生成嵌入向量失败: {embed_error}")
                        continue

        print(f"总共创建了 {len(all_texts)} 个有效节点")
        if not all_texts:
            gr.Info("没有有效的文本节点")
            return

        # 用pymilvus原生操作插入数据
        print("开始用pymilvus原生操作插入数据...")
        try:
            from pymilvus import connections, Collection, utility
            
            # 连接到Milvus
            connections.connect(
                alias="default",
                host=MILVUS_HOST,
                port=MILVUS_PORT,
                user=MILVUS_USER if MILVUS_USER else None,
                password=MILVUS_PASSWORD if MILVUS_PASSWORD else None
            )
            
            # 获取集合
            collection = Collection(db_name)
            
            # 批量插入数据
            insert_data = [
                all_texts,
                all_embeddings
            ]
            collection.insert(insert_data)
            print(f"✅ 数据插入成功，插入 {len(all_texts)} 条记录")
            
            # 刷新集合
            collection.flush()
            print("✅ 集合刷新成功")
            
            # 加载集合
            collection.load()
            print("✅ 集合加载成功")
            
            # 验证数据
            num_entities = collection.num_entities
            print(f"✅ 集合包含 {num_entities} 个实体")
            
            connections.disconnect("default")
            
            # 最终验证
            if num_entities > 0:
                gr.Info(f"Milvus知识库创建成功！集合名称: {db_name}，存储了 {num_entities} 个向量，可前往RAG问答进行提问")
            else:
                gr.Info(f"Milvus知识库创建完成，集合 {db_name} 存在但未检测到数据，请检查文档内容")
            
        except Exception as e:
            print(f"❌ 数据插入失败: {e}")
            import traceback
            traceback.print_exc()
            gr.Info("数据插入失败")
            return

    except Exception as e:
        print(f"手动创建Milvus知识库失败: {e}")
        import traceback
        traceback.print_exc()
        gr.Info(f"手动创建Milvus知识库失败: {str(e)}")

def create_structured_milvus_db(db_name: str, data_table: list):
    ensure_event_loop()
    """创建结构化数据的Milvus知识库"""
    print(f"Milvus知识库名称为：{db_name}，数据表名称为：{data_table}")
    
    if data_table is None:
        gr.Info("没有选择数据表")
        return
    elif len(db_name) == 0:
        gr.Info("没有命名知识库")
        return
    
    # 检查Milvus集合是否已存在
    existing_collections = get_milvus_collections()
    
    # 检查原始名称和规范化名称
    normalized_name = normalize_collection_name(db_name)
    if db_name in existing_collections or normalized_name in existing_collections:
        gr.Info("Milvus知识库已存在，请换个名字或删除原来知识库再创建")
        return
    
    try:
        gr.Info("正在创建结构化Milvus知识库...")
        
        # 加载文档
        documents = []
        for label in data_table:
            label_path = os.path.join(STRUCTURED_FILE_PATH, label)
            print(f"检查路径: {label_path}, 存在: {os.path.exists(label_path)}")
            if os.path.exists(label_path):
                docs = SimpleDirectoryReader(label_path).load_data()
                print(f"从 {label} 加载了 {len(docs)} 个文档")
                documents.extend(docs)
        
        print(f"总共加载了 {len(documents)} 个文档")
        if not documents:
            gr.Info("没有找到文档数据")
            return
        
        # 创建Milvus向量存储
        print(f"开始创建Milvus向量存储，集合名称: {db_name}")
        vector_store = create_milvus_vector_store(db_name)
        if vector_store is None:
            gr.Info("创建Milvus向量存储失败")
            return
        
        # 使用LlamaIndex创建索引
        print("开始使用LlamaIndex创建索引...")
        index = VectorStoreIndex.from_documents(
            documents,
            vector_store=vector_store
        )
        
        print("结构化Milvus知识库创建完成，验证数据存储...")
        
        # 验证数据是否成功存储
        try:
            from pymilvus import connections, utility
            connections.connect(
                alias="default",
                host=MILVUS_HOST,
                port=MILVUS_PORT,
                user=MILVUS_USER if MILVUS_USER else None,
                password=MILVUS_PASSWORD if MILVUS_PASSWORD else None
            )
            
            # 获取所有集合名称
            all_collections = utility.list_collections()
            print(f"Milvus中的所有集合: {all_collections}")
            
            # 检查原始名称和规范化名称
            found_collection = None
            if utility.has_collection(db_name):
                found_collection = db_name
                print(f"找到原始集合名称: {db_name}")
            elif utility.has_collection(normalized_name):
                found_collection = normalized_name
                print(f"找到规范化集合名称: {normalized_name}")
            else:
                # 检查是否有其他可能的集合名称
                for collection_name in all_collections:
                    if db_name.lower() in collection_name.lower() or normalized_name.lower() in collection_name.lower():
                        found_collection = collection_name
                        print(f"找到相似集合名称: {collection_name}")
                        break
            
            if found_collection:
                collection = utility.get_collection(found_collection)
                num_entities = collection.num_entities
                print(f"集合 {found_collection} 存在，包含 {num_entities} 个实体")
                if num_entities > 0:
                    gr.Info(f"结构化Milvus知识库创建成功！集合名称: {found_collection}，存储了 {num_entities} 个向量，可前往RAG问答进行提问")
                else:
                    gr.Info(f"结构化Milvus知识库创建完成，集合 {found_collection} 存在但未检测到数据，请检查文档内容")
            else:
                print(f"未找到集合 {db_name} 或 {normalized_name}")
                print(f"可用的集合: {all_collections}")
                gr.Info("结构化Milvus知识库创建失败，集合未找到")
            
            connections.disconnect("default")
            
        except Exception as verify_error:
            print(f"验证数据存储时出错: {verify_error}")
            import traceback
            traceback.print_exc()
            gr.Info("结构化Milvus知识库创建完成，但验证失败")
        
    except Exception as e:
        print(f"创建结构化Milvus知识库失败: {e}")
        import traceback
        traceback.print_exc()
        gr.Info(f"创建结构化Milvus知识库失败: {str(e)}")

def delete_milvus_db(db_name: str):
    ensure_event_loop()
    """删除Milvus知识库"""
    if db_name is None:
        gr.Info("请选择要删除的知识库")
        return
    
    try:
        from pymilvus import connections, utility
        
        # 连接到Milvus
        connections.connect(
            alias="default",
            host=MILVUS_HOST,
            port=MILVUS_PORT,
            user=MILVUS_USER if MILVUS_USER else None,
            password=MILVUS_PASSWORD if MILVUS_PASSWORD else None
        )
        
        # 检查原始名称和规范化名称
        normalized_name = normalize_collection_name(db_name)
        deleted = False
        
        if utility.has_collection(db_name):
            utility.drop_collection(db_name)
            gr.Info(f"已成功删除Milvus知识库: {db_name}")
            print(f"已成功删除Milvus知识库: {db_name}")
            deleted = True
        elif utility.has_collection(normalized_name):
            utility.drop_collection(normalized_name)
            gr.Info(f"已成功删除Milvus知识库: {normalized_name}")
            print(f"已成功删除Milvus知识库: {normalized_name}")
            deleted = True
        else:
            gr.Info(f"Milvus知识库 {db_name} 或 {normalized_name} 不存在")
            print(f"Milvus知识库 {db_name} 或 {normalized_name} 不存在")
        
        connections.disconnect("default")
        
    except Exception as e:
        print(f"删除Milvus知识库失败: {e}")
        gr.Info(f"删除Milvus知识库失败: {str(e)}")

def update_milvus_knowledge_base():
    ensure_event_loop()
    """更新Milvus知识库列表"""
    collections = get_milvus_collections()
    return gr.update(choices=collections)

def create_tmp_milvus_kb(files):
    ensure_event_loop()
    """创建临时Milvus知识库"""
    try:
        if not os.path.exists(os.path.join("File", TMP_NAME)):
            os.mkdir(os.path.join("File", TMP_NAME))
        
        # 移动文件到临时目录
        for file in files:
            file_name = os.path.basename(file)
            import shutil
            shutil.move(file, os.path.join("File", TMP_NAME, file_name))
        
        # 加载文档
        documents = SimpleDirectoryReader(os.path.join("File", TMP_NAME)).load_data()
        
        if not documents:
            print("没有找到文档数据")
            return
        
        print(f"从临时目录加载了 {len(documents)} 个文档")
        
        # 使用抽取的方法创建临时Milvus集合
        collection = create_milvus_collection(TMP_NAME)
        if collection is None:
            print("创建临时Milvus集合失败")
            return
        
        print("临时Milvus向量存储创建成功，开始处理文档...")
        
        from llama_index.core.schema import TextNode
        from llama_index.core.node_parser import SentenceSplitter
        
        text_splitter = SentenceSplitter(chunk_size=512, chunk_overlap=50)
        all_texts = []
        all_embeddings = []
        
        for doc in documents:
            text_chunks = text_splitter.split_text(doc.get_content())
            print(f"文档 {doc.get_doc_id()} 分割为 {len(text_chunks)} 个块")
            for i, chunk in enumerate(text_chunks):
                if chunk.strip() and len(chunk.strip()) >= 10:
                    try:
                        embedding = EMBED_MODEL.get_text_embedding(chunk)
                        if embedding is None or len(embedding) == 0:
                            print(f"警告：节点 {i} 的嵌入向量为空，跳过")
                            continue
                        print(f"节点 {i} 嵌入向量维度: {len(embedding)}")
                        all_texts.append(chunk)
                        all_embeddings.append(embedding)
                    except Exception as embed_error:
                        print(f"警告：节点 {i} 生成嵌入向量失败: {embed_error}")
                        continue
        
        print(f"总共创建了 {len(all_texts)} 个有效节点")
        if not all_texts:
            print("没有有效的文本节点")
            return
        
        # 用pymilvus原生操作插入数据
        print("开始用pymilvus原生操作插入临时数据...")
        try:
            from pymilvus import connections, Collection, utility
            
            # 连接到Milvus
            connections.connect(
                alias="default",
                host=MILVUS_HOST,
                port=MILVUS_PORT,
                user=MILVUS_USER if MILVUS_USER else None,
                password=MILVUS_PASSWORD if MILVUS_PASSWORD else None
            )
            
            # 获取集合
            collection = Collection(TMP_NAME)
            
            # 批量插入数据
            insert_data = [
                all_texts,
                all_embeddings
            ]
            collection.insert(insert_data)
            print(f"✅ 临时数据插入成功，插入 {len(all_texts)} 条记录")
            
            # 刷新集合
            collection.flush()
            print("✅ 临时集合刷新成功")
            
            # 加载集合
            collection.load()
            print("✅ 临时集合加载成功")
            
            # 验证数据
            num_entities = collection.num_entities
            print(f"✅ 临时集合包含 {num_entities} 个实体")
            
            connections.disconnect("default")
            
            # 最终验证
            if num_entities > 0:
                print(f"临时Milvus知识库创建成功！集合名称: {TMP_NAME}，存储了 {num_entities} 个向量")
            else:
                print(f"临时Milvus知识库创建完成，集合 {TMP_NAME} 存在但未检测到数据，请检查文档内容")
                
        except Exception as e:
            print(f"插入临时数据失败: {e}")
            import traceback
            traceback.print_exc()
        
    except Exception as e:
        print(f"创建临时Milvus知识库失败: {e}")
        import traceback
        traceback.print_exc()

def clear_tmp_milvus():
    ensure_event_loop()
    """清除临时Milvus知识库"""
    try:
        # 删除临时文件目录
        if os.path.exists(os.path.join("File", TMP_NAME)):
            import shutil
            shutil.rmtree(os.path.join("File", TMP_NAME))
        
        # 删除临时Milvus集合
        from pymilvus import connections, utility
        
        connections.connect(
            alias="default",
            host=MILVUS_HOST,
            port=MILVUS_PORT,
            user=MILVUS_USER if MILVUS_USER else None,
            password=MILVUS_PASSWORD if MILVUS_PASSWORD else None
        )
        
        # 检查原始名称和规范化名称
        normalized_tmp_name = normalize_collection_name(TMP_NAME)
        
        if utility.has_collection(TMP_NAME):
            utility.drop_collection(TMP_NAME)
            print(f"已删除临时集合: {TMP_NAME}")
        elif utility.has_collection(normalized_tmp_name):
            utility.drop_collection(normalized_tmp_name)
            print(f"已删除临时集合: {normalized_tmp_name}")
        
        connections.disconnect("default")
        
    except Exception as e:
        print(f"清除临时Milvus知识库失败: {e}")
